package _01_第一季.树04.Tree02;

import java.util.Comparator;

public class AVLTree<E> extends BBST<E>{
    public AVLTree(){
        this(null);
    }

    public AVLTree(Comparator<E> comparator){
        super(comparator);
    }

    private static class AVLNode<E> extends Node<E>{
        int height = 1;
        public AVLNode(){
            this(null,null);
        }
        public AVLNode(E element,Node<E> parent){
            super(element, parent);
        }

        public int balanceFactor(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = 1 + Math.max(leftHeight,rightHeight);
        }

        public Node<E> tallerChild(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            if(leftHeight>rightHeight) return left;
            if(leftHeight<rightHeight) return right;
            return isLeftChild() ? left : right;
        }

        @Override
        public String toString() {
            String parentString = "null";
            if(parent!=null){
                parentString = parent.element.toString();
            }
            return element+"_"+parentString+"_"+height;
        }
    }

    //添加之后的修复
    @Override
    protected void afterAdd(Node<E> node) {
       while((node = node.parent)!=null){
           if(isBalanced(node)){
               updateHeight(node); //更新高度
           }else{
                reBalance(node);
                break;
           }
       }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while((node = node.parent)!=null){
            if(isBalanced(node)){
                updateHeight(node); //更新高度
            }else{
                reBalance(node);
            }
        }
    }

    //修复平衡，判断方向
    private void reBalance(Node<E> grand){
        Node<E> parent = ((AVLNode<E>)grand).tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        if(parent.isLeftChild()){
            if(node.isLeftChild()){
                rotateRight(grand);
            }else{ //LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else{
            if(node.isLeftChild()){ //RL
                rotateRight(parent);
                rotateLeft(grand);
            }else{ //RR
                rotateLeft(grand);
            }
        }
    }

    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);

        updateHeight(grand);
        updateHeight(parent);
    }

    //更新节点的方法
    public void updateHeight(Node<E> node){
        ((AVLNode<E>)node).updateHeight();
    }

    //判断是否平衡
    public boolean isBalanced(Node<E> node){
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element,parent);
    }
}
